##计算ATR数值
from typing import Optional, Dict
from datetime import datetime, timedelta
from collections import deque
import time
import pandas as pd
import numpy as np
import talib as ta
from 主程.基础配置 import settings
from xtquant import xtdata


class ATRCalculator:
    """
    ATR 波动率计算器（基于当日 tick 重采样为 1 分钟K）

    - 使用 True Range 定义：TR = max(high-low, |high-prev_close|, |low-prev_close|)
    - ATR 计算方式：默认 TA-Lib（Wilder），可选 EMA/SMA
    - 内置简单缓存，避免每个tick都重复重算
    """

    def __init__(self, data_provider=xtdata, lookback: int = 14, method: str = "talib",
                 cache_ttl_seconds: int = 10, verbose: bool = False):
        self._xt = data_provider
        self.lookback = int(lookback)
        self.method = method.lower().strip()  # "talib" | "ema" | "sma"
        self.cache_ttl_seconds = int(cache_ttl_seconds)
        self.verbose = bool(verbose)
        # 缓存: code -> {"ts": float, "atr": float}
        self._cache: Dict[str, Dict[str, float]] = {}

    def _log(self, msg: str):
        if self.verbose:
            print(f"【ATR】{msg}")

    def _now(self) -> float:
        return time.time()

    def _download_today_ticks(self, stock_code: str) -> Optional[pd.DataFrame]:
        """
        下载并读取今日 tick 数据
        """
        try:
            today = datetime.now().strftime("%Y%m%d")
            self._xt.download_history_data(stock_code=stock_code, period="tick", start_time=today, end_time=today)
            time.sleep(1.5)
            data = self._xt.get_local_data(stock_list=[stock_code], period='tick', start_time=today, end_time=today)
            if data and stock_code in data:
                return data[stock_code]
            return None
        except Exception:
            return None

    def _ensure_time_and_price(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
        if df is None or df.empty:
            return None
        price_col = next((col for col in ['price', 'lastPrice', 'close', 'last', 'current'] if col in df.columns), None)
        if not price_col:
            return None
        time_col = next((col for col in df.columns if 'time' in col.lower()), None)
        if not time_col:
            return None
        try:
            if df[time_col].dtype in ['int64', 'float64']:
                df['time'] = pd.to_datetime(df[time_col], unit='ms') + pd.Timedelta(hours=8)
            else:
                df['time'] = pd.to_datetime(df[time_col], format='%Y%m%d%H%M%S', errors='coerce')
            df = df[df['time'].notna()].sort_values('time')
        except Exception:
            return None
        if df.empty:
            return None
        # 过滤价格<=0 及午休
        df = df[df[price_col] > 0].copy()
        lunch_start = pd.Timestamp('11:30:00').time()
        lunch_end = pd.Timestamp('12:59:59').time()
        df['time_only'] = df['time'].dt.time
        df = df[~((df['time_only'] >= lunch_start) & (df['time_only'] <= lunch_end))]
        if df.empty:
            return None
        df.rename(columns={price_col: 'price'}, inplace=True)
        return df

    def _resample_to_minutes(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
        try:
            df = df.set_index('time')
            ohlc = df['price'].resample('1min').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
            }).dropna()

            # 仅保留最后 lookback 根K线（确保ATR计算范围准确）
            ohlc = ohlc.iloc[-self.lookback:]

            # 打印调试信息
            if not ohlc.empty:
                print(f"K线时间范围: {ohlc.index[0]} 至 {ohlc.index[-1]}")
                print(f"K线数量: {len(ohlc)}")

            return ohlc.reset_index()
        except Exception as e:
            print(f"重采样失败: {e}")
            return None


    def _compute_tr(self, df_minutes: pd.DataFrame) -> Optional[pd.Series]:
        if df_minutes is None or df_minutes.empty:
            return None
        high = df_minutes['high'].astype(float)
        low = df_minutes['low'].astype(float)
        close = df_minutes['close'].astype(float)
        prev_close = close.shift(1)
        tr = pd.concat([
            (high - low),
            (high - prev_close).abs(),
            (low - prev_close).abs()
        ], axis=1).max(axis=1)
        return tr

    def _ema(self, series: pd.Series, period: int) -> pd.Series:
        if period <= 0 or series is None or series.empty:
            return pd.Series(dtype=float)
        alpha = 2 / (period + 1)
        ema_vals = []
        last = None
        for val in series:
            if last is None:
                last = val
            else:
                last = alpha * val + (1 - alpha) * last
            ema_vals.append(last)
        return pd.Series(ema_vals, index=series.index)

    def _sma(self, series: pd.Series, period: int) -> pd.Series:
        try:
            return series.rolling(window=period, min_periods=period).mean()
        except Exception:
            return pd.Series(dtype=float)

    def _compute_atr(self, df_minutes: pd.DataFrame, lookback: int) -> Optional[float]:
        tr = self._compute_tr(df_minutes)
        if tr is None or tr.empty:
            return None

        # 获取最新价格用于异常检查
        current_price = df_minutes['close'].iloc[-1] if not df_minutes.empty else None

        # 修复：仅使用最后lookback根K线
        tr = tr.iloc[-lookback:]
        if len(tr) < lookback:
            return None

        # 添加异常波动过滤（使用当前价格检查）
        if current_price is not None and any(tr > current_price * 0.1):  # 过滤单根K线波动>10%的情况
            return None

        method_lower = self.method.lower()
        if method_lower == 'talib':
            # 使用 TA-Lib 的 ATR（Wilder 平滑）
            high = df_minutes['high'].astype(float).values
            low = df_minutes['low'].astype(float).values
            close = df_minutes['close'].astype(float).values
            try:
                atr_values = ta.ATR(high, low, close, timeperiod=lookback)
            except Exception:
                return None
            if atr_values is None or len(atr_values) == 0 or np.isnan(atr_values[-1]):
                return None
            return float(atr_values[-1])

        if method_lower == 'ema':
            atr_series = self._ema(tr, lookback)
        else:
            atr_series = self._sma(tr, lookback)

        return float(atr_series.iloc[-1]) if not atr_series.empty else None

    def _filter_data_before_time(self, df: pd.DataFrame, end_time: datetime) -> Optional[pd.DataFrame]:
        """
        过滤掉指定时间之后的数据（回测模式使用）
        
        Args:
            df: 原始数据DataFrame
            end_time: 结束时间
            
        Returns:
            过滤后的DataFrame，如果过滤失败则返回None
        """
        try:
            if df.empty:
                return df
            
            # 确保时间列存在
            if 'time' not in df.columns:
                return df
            
            # 过滤掉结束时间之后的数据
            filtered_df = df[df['time'] <= end_time]
            return filtered_df
            
        except Exception as e:
            self._log(f"数据过滤异常: {e}")
            return None

    def calculate_atr(self, stock_code: str, current_price: Optional[float] = None, lookback: Optional[int] = None, end_time: Optional[datetime] = None) -> Optional[float]:
        """
        返回最近的 ATR 数值（按 1 分钟K 计算）。内置缓存，默认 10 秒刷新一次。
        
        Args:
            stock_code: 股票代码
            current_price: 当前价格
            lookback: 回看期数
            end_time: 结束时间（回测模式使用，实时模式为None）
        """
        lookback = int(lookback or self.lookback)

        # 缓存命中（回测模式下不使用缓存）
        if end_time is None:
            cached = self._cache.get(stock_code)
            if cached and (self._now() - cached.get('ts', 0)) < self.cache_ttl_seconds:
                return cached.get('atr')

        # 下载并处理当日 tick
        ticks_df = self._download_today_ticks(stock_code)
        ticks_df = self._ensure_time_and_price(ticks_df)
        if ticks_df is None:
            self._log(f"{stock_code} 无有效tick数据，ATR返回None")
            return None
            
        # 在回测模式下，过滤掉结束时间之后的数据
        if end_time is not None:
            ticks_df = self._filter_data_before_time(ticks_df, end_time)
            if ticks_df is None or ticks_df.empty:
                self._log(f"{stock_code} 回测模式下过滤后无有效数据，ATR返回None")
                return None
        
        minutes_df = self._resample_to_minutes(ticks_df)
        if minutes_df is None or minutes_df.empty:
            self._log(f"{stock_code} 分钟数据为空，ATR返回None")
            return None

        atr_value = self._compute_atr(minutes_df, lookback)
        if atr_value is None:
            self._log(f"{stock_code} ATR计算失败")
            return None

        # 更新缓存（回测模式下不更新缓存）
        if end_time is None:
            self._cache[stock_code] = {"ts": self._now(), "atr": atr_value}
            self._log(f"{stock_code} ATR={atr_value:.6f} (lookback={lookback},{self.method.upper()})")
            print(f"{stock_code} ATR={atr_value:.6f} (lookback={lookback},{self.method.upper()})")

        # 输出详细评估信息
        try:
            last_price = float(current_price) if (current_price is not None) else float(minutes_df['close'].iloc[-1])
            if last_price > 0:
                atr_ratio = atr_value / last_price
                threshold = max(0.0006, atr_ratio * 0.7)
                need_buy = atr_ratio > threshold
                print(f"ATR绝对值: {atr_value:.6f}")
                print(f"ATR比率: {atr_ratio:.4%}")
                print(f"阈值: {threshold:.4%}")
                print(f"是否买入: {need_buy}")
        except Exception:
            pass

        return atr_value


# 用法示例（单独运行调试）
if __name__ == "__main__":
    calc = ATRCalculator(xtdata, lookback=14, method='talib')  # 使用 TA-Lib 计算 ATR
    code = "513060.SH"
    atr = calc.calculate_atr(code)
    last_price = 0.722
    need_buy = False

    if atr is not None:
        atr_ratio = atr / last_price
        ##threshold = 0.001 * settings.ATR_X
        threshold = max(0.0006, atr_ratio * 0.7)  # 动态阈值
        if atr_ratio > threshold:
            need_buy = True

        print(f"ATR绝对值: {atr:.6f}")
        print(f"ATR比率: {atr_ratio:.4%}")  # 格式化为百分比
        print(f"阈值: {threshold:.4%}")
        print(f"是否买入: {need_buy}")